package main

import (
	"unicode"
)

// TokenType represents the type of a token
type TokenType int

const (
	TokenEOF TokenType = iota
	TokenLeftParen
	TokenRightParen
	TokenDot
	TokenNumber
	TokenString
	TokenSymbol
	TokenQuote
	TokenBoolean
	TokenChar
	TokenVector
)

// Token represents a lexical token
type Token struct {
	Type  TokenType
	Value string
	Line  int
	Col   int
}

// Tokenizer represents the lexical analyzer
type Tokenizer struct {
	runes   []rune
	pos     int
	current rune
	line    int
	col     int
}

// NewTokenizer creates a new tokenizer
func NewTokenizer(input string) *Tokenizer {
	runes := []rune(input)
	t := &Tokenizer{
		runes: runes,
		pos:   0,
		line:  1,
		col:   1,
	}
	if len(runes) > 0 {
		t.current = runes[0]
	} else {
		t.current = 0
	}
	return t
}

// advance moves to the next character
func (t *Tokenizer) advance() {
	if t.pos < len(t.runes)-1 {
		if t.current == '\n' {
			t.line++
			t.col = 1
		} else {
			t.col++
		}
		t.pos++
		t.current = t.runes[t.pos]
	} else {
		t.current = 0 // EOF
	}
}

// peek returns the next character without advancing
func (t *Tokenizer) peek() rune {
	if t.pos < len(t.runes)-1 {
		return t.runes[t.pos+1]
	}
	return 0
}

// skipWhitespace skips whitespace and comments
func (t *Tokenizer) skipWhitespace() {
	for t.current != 0 && (unicode.IsSpace(t.current) || t.current == ';') {
		if t.current == ';' {
			// Skip comment until end of line
			for t.current != 0 && t.current != '\n' {
				t.advance()
			}
		} else {
			t.advance()
		}
	}
}

// readString reads a string literal
func (t *Tokenizer) readString() string {
	var result []rune
	t.advance() // skip opening quote

	for t.current != 0 && t.current != '"' {
		if t.current == '\\' {
			t.advance()
			switch t.current {
			case 'n':
				result = append(result, '\n')
			case 't':
				result = append(result, '\t')
			case 'r':
				result = append(result, '\r')
			case '\\':
				result = append(result, '\\')
			case '"':
				result = append(result, '"')
			default:
				result = append(result, t.current)
			}
		} else {
			result = append(result, t.current)
		}
		t.advance()
	}

	if t.current == '"' {
		t.advance() // skip closing quote
	}

	return string(result)
}

// readNumber reads a number (integer, rational, real, or complex)
// Supports formats like: 123, -456, 1/2, 3.14, 1e10, #e3.14, #i22/7, 3+4i, -2i, +i
func (t *Tokenizer) readNumber() string {
	var result string

	// Handle exactness prefix (#e or #i)
	if t.current == '#' && (t.peek() == 'e' || t.peek() == 'i') {
		result += string(t.current) // #
		t.advance()
		result += string(t.current) // e or i
		t.advance()
	}

	// Read the first part (real part or complete number)
	firstPart := t.readRealPart()
	result += firstPart

	// Check for imaginary unit 'i' immediately after the first part
	if t.current == 'i' {
		// This is a pure imaginary number (like 5i or -3i)
		result += "i"
		t.advance()
		return result
	}

	// Check for complex number notation (real+imagi or real-imagi)
	if t.current == '+' || t.current == '-' {
		// Save position in case this isn't actually a complex number
		complexPos := t.pos
		complexCurrent := t.current
		complexLine := t.line
		complexCol := t.col

		// Read the imaginary part
		imagSign := string(t.current)
		t.advance()

		// Handle pure 'i' (like 3+i or 3-i)
		if t.current == 'i' {
			result += imagSign + "i"
			t.advance()
			return result
		}

		// Read the imaginary coefficient
		imagPart := t.readRealPart()
		if imagPart != "" && t.current == 'i' {
			// Valid complex number
			result += imagSign + imagPart + "i"
			t.advance()
			return result
		} else {
			// Not a complex number, backtrack
			t.pos = complexPos
			t.current = complexCurrent
			t.line = complexLine
			t.col = complexCol
		}
	}

	return result
}

// readRealPart reads a real number part (without exactness prefix)
func (t *Tokenizer) readRealPart() string {
	var result string

	// Handle sign
	if t.current == '-' || t.current == '+' {
		result += string(t.current)
		t.advance()
	}

	// Read digits before decimal point or slash
	for t.current != 0 && unicode.IsDigit(t.current) {
		result += string(t.current)
		t.advance()
	}

	// Handle rational numbers (1/2)
	if t.current == '/' {
		result += string(t.current)
		t.advance()
		// Read denominator
		for t.current != 0 && unicode.IsDigit(t.current) {
			result += string(t.current)
			t.advance()
		}
		return result
	}

	// Handle decimal point
	if t.current == '.' {
		result += string(t.current)
		t.advance()
		// Read digits after decimal point
		for t.current != 0 && unicode.IsDigit(t.current) {
			result += string(t.current)
			t.advance()
		}
	}

	// Handle scientific notation (e or E)
	if t.current == 'e' || t.current == 'E' {
		result += string(t.current)
		t.advance()
		// Handle sign in exponent
		if t.current == '-' || t.current == '+' {
			result += string(t.current)
			t.advance()
		}
		// Read exponent digits
		for t.current != 0 && unicode.IsDigit(t.current) {
			result += string(t.current)
			t.advance()
		}
	}

	return result
}

// readChar reads a character literal
func (t *Tokenizer) readChar() string {
	var result string
	t.advance() // skip #
	t.advance() // skip \\

	// Handle named characters
	if unicode.IsLetter(t.current) {
		// Read character name
		var name string
		for t.current != 0 && (unicode.IsLetter(t.current) || unicode.IsDigit(t.current)) {
			name += string(t.current)
			t.advance()
		}
		// Check for named characters
		switch name {
		case "space":
			return "#\\space"
		case "tab":
			return "#\\tab"
		case "newline":
			return "#\\newline"
		case "return":
			return "#\\return"
		default:
			// Single character name (like "a", "A", etc.)
			if len(name) == 1 {
				return "#\\" + name
			}
			// Invalid character name
			return "#\\" + name
		}
	} else {
		// Single character
		result = "#\\" + string(t.current)
		t.advance()
		return result
	}
}

// readSymbol reads a symbol or identifier
func (t *Tokenizer) readSymbol() string {
	var result string

	for t.current != 0 && !unicode.IsSpace(t.current) &&
		t.current != '(' && t.current != ')' && t.current != '"' && t.current != '\'' {
		result += string(t.current)
		t.advance()
	}

	return result
}

// isNumber checks if the current position starts a valid number
// This needs to distinguish between pure numbers and identifiers starting with digits
func (t *Tokenizer) isNumber() bool {
	// Save current position
	savedPos := t.pos
	savedCurrent := t.current
	savedLine := t.line
	savedCol := t.col

	// Check for exactness prefix
	if t.current == '#' && (t.peek() == 'e' || t.peek() == 'i') {
		t.advance() // skip #
		t.advance() // skip e or i
	}

	// Check for sign
	if t.current == '-' || t.current == '+' {
		t.advance()
		if !unicode.IsDigit(t.current) && t.current != 'i' {
			// Restore position and return false
			t.pos = savedPos
			t.current = savedCurrent
			t.line = savedLine
			t.col = savedCol
			return false
		}
	} else if !unicode.IsDigit(t.current) && t.current != 'i' {
		// Restore position and return false
		t.pos = savedPos
		t.current = savedCurrent
		t.line = savedLine
		t.col = savedCol
		return false
	}

	// Check for pure imaginary (i, +i, -i)
	if t.current == 'i' {
		// Look ahead to see if it's the end of the token
		next := t.peek()
		isValidEnd := next == 0 || unicode.IsSpace(next) || next == '(' || next == ')' || next == '\'' || next == '"'
		if isValidEnd {
			// Restore position and return true for pure imaginary
			t.pos = savedPos
			t.current = savedCurrent
			t.line = savedLine
			t.col = savedCol
			return true
		}
	}

	// Read through digits
	for t.current != 0 && unicode.IsDigit(t.current) {
		t.advance()
	}

	// Check what comes after the digits
	isValidNumber := false
	if t.current == 0 || unicode.IsSpace(t.current) || t.current == '(' || t.current == ')' || t.current == '\'' {
		// End of token or delimiter - this is a pure number
		isValidNumber = true
	} else if t.current == '.' || t.current == '/' || t.current == 'e' || t.current == 'E' {
		// Decimal, rational, or scientific notation - this is a number
		isValidNumber = true
	} else if t.current == '+' || t.current == '-' {
		// Might be a complex number like 3+4i or 3-2i
		// Look ahead for imaginary part
		t.advance() // skip + or -
		// Check for pure i (like 3+i)
		if t.current == 'i' {
			next := t.peek()
			isValidEnd := next == 0 || unicode.IsSpace(next) || next == '(' || next == ')' || next == '\'' || next == '"'
			if isValidEnd {
				isValidNumber = true
			}
		} else {
			// Read digits for imaginary coefficient
			hasDigits := false
			for t.current != 0 && unicode.IsDigit(t.current) {
				t.advance()
				hasDigits = true
			}
			// Check for 'i' after digits
			if hasDigits && t.current == 'i' {
				next := t.peek()
				isValidEnd := next == 0 || unicode.IsSpace(next) || next == '(' || next == ')' || next == '\'' || next == '"'
				if isValidEnd {
					isValidNumber = true
				}
			}
		}
	} else if t.current == 'i' {
		// Check for 'i' suffix (pure imaginary like 5i)
		next := t.peek()
		isValidEnd := next == 0 || unicode.IsSpace(next) || next == '(' || next == ')' || next == '\'' || next == '"'
		if isValidEnd {
			isValidNumber = true
		}
	} else {
		// Has non-numeric characters after digits - this is an identifier
		isValidNumber = false
	}

	// Restore position
	t.pos = savedPos
	t.current = savedCurrent
	t.line = savedLine
	t.col = savedCol

	return isValidNumber
}

// NextToken returns the next token
func (t *Tokenizer) NextToken() *Token {
	for {
		t.skipWhitespace()

		if t.current == 0 {
			return &Token{Type: TokenEOF, Line: t.line, Col: t.col}
		}

		switch t.current {
		case '(':
			token := &Token{Type: TokenLeftParen, Value: "(", Line: t.line, Col: t.col}
			t.advance()
			return token

		case ')':
			token := &Token{Type: TokenRightParen, Value: ")", Line: t.line, Col: t.col}
			t.advance()
			return token

		case '\'':
			token := &Token{Type: TokenQuote, Value: "'", Line: t.line, Col: t.col}
			t.advance()
			return token

		case '"':
			value := t.readString()
			return &Token{Type: TokenString, Value: value, Line: t.line, Col: t.col}

		case '#':
			if t.peek() == 't' || t.peek() == 'f' {
				line, col := t.line, t.col
				t.advance() // skip #
				value := "#" + string(t.current)
				t.advance()
				return &Token{Type: TokenBoolean, Value: value, Line: line, Col: col}
			} else if t.peek() == '\\' {
				// Character literal
				line, col := t.line, t.col
				value := t.readChar()
				return &Token{Type: TokenChar, Value: value, Line: line, Col: col}
			} else if t.peek() == 'e' || t.peek() == 'i' {
				// Exactness prefix
				value := t.readNumber()
				return &Token{Type: TokenNumber, Value: value, Line: t.line, Col: t.col}
			} else if t.peek() == '(' {
				// Vector literal
				line, col := t.line, t.col
				// Skip # and (
				t.advance()
				t.advance()
				// Read the vector content until closing )
				var vectorContent []rune
				depth := 1
				for t.current != 0 && depth > 0 {
					vectorContent = append(vectorContent, t.current)
					if t.current == '(' {
						depth++
					} else if t.current == ')' {
						depth--
					}
					t.advance()
				}
				// Remove the last ) from the content
				if len(vectorContent) > 0 && vectorContent[len(vectorContent)-1] == ')' {
					vectorContent = vectorContent[:len(vectorContent)-1]
				}
				value := "#(" + string(vectorContent) + ")"
				return &Token{Type: TokenVector, Value: value, Line: line, Col: col}
			} else {
				// Unknown # construct, treat as symbol
				value := t.readSymbol()
				return &Token{Type: TokenSymbol, Value: value, Line: t.line, Col: t.col}
			}

		case '.':
			// Check if this is part of a number (decimal point)
			if unicode.IsDigit(t.peek()) {
				// This is a decimal number, not a dot token
				value := t.readNumber()
				return &Token{Type: TokenNumber, Value: value, Line: t.line, Col: t.col}
			} else {
				// This is a dot token for dotted pairs
				token := &Token{Type: TokenDot, Value: ".", Line: t.line, Col: t.col}
				t.advance()
				return token
			}

		default:
			// Check for pure imaginary number (starting with i, +i, or -i)
			if t.current == 'i' && (t.peek() == 0 || unicode.IsSpace(t.peek()) || t.peek() == '(' || t.peek() == ')') {
				// Pure 'i'
				value := "i"
				t.advance()
				return &Token{Type: TokenNumber, Value: value, Line: t.line, Col: t.col}
			} else if (t.current == '+' || t.current == '-') && t.peek() == 'i' {
				// Check if this is '+i' or '-i'
				savedPos := t.pos
				savedCurrent := t.current
				savedLine := t.line
				savedCol := t.col

				line, col := t.line, t.col
				sign := string(t.current)
				t.advance() // move to 'i'
				if t.current == 'i' {
					t.advance() // move past 'i'
					// Check if this is the end of the token
					if t.current == 0 || unicode.IsSpace(t.current) || t.current == '(' || t.current == ')' {
						return &Token{Type: TokenNumber, Value: sign + "i", Line: line, Col: col}
					}
				}
				// Restore position - this isn't a pure imaginary number
				t.pos = savedPos
				t.current = savedCurrent
				t.line = savedLine
				t.col = savedCol
			}

			// Check for number
			if t.isNumber() {
				value := t.readNumber()
				return &Token{Type: TokenNumber, Value: value, Line: t.line, Col: t.col}
			} else {
				value := t.readSymbol()
				return &Token{Type: TokenSymbol, Value: value, Line: t.line, Col: t.col}
			}
		}
	}
}

// Tokenize returns all tokens from the input
func Tokenize(input string) ([]*Token, error) {
	tokenizer := NewTokenizer(input)
	var tokens []*Token

	for {
		token := tokenizer.NextToken()
		tokens = append(tokens, token)
		if token.Type == TokenEOF {
			break
		}
	}

	return tokens, nil
}
